# -*- coding: utf-8 -*-
"""
文件名: yasa_bandpower_helper.py
用途: 使用YASA的bandpower方法优化频域比例生成，专门为前端显示提供频域比例数据
方法: 
- 使用YASA的bandpower函数计算各频段功率
- 支持从MNE Raw对象或numpy数组计算频段功率
- 提供标准化的频段比例输出格式
原理:
- YASA的bandpower基于Welch方法计算功率谱密度
- 自动处理频段定义和功率积分
- 比手动实现更准确和高效
注意:
- 这个功能仅用于前端显示，没有其他实际作用
- 需要安装yasa库: pip install yasa
- 输入数据应为微伏单位的EEG信号
"""

import numpy as np
import pandas as pd
import mne
import warnings
from typing import Dict, List, Optional, Union, Tuple
import logging

logger = logging.getLogger(__name__)

# 尝试导入YASA
try:
    import yasa
    YASA_AVAILABLE = True
    logger.info("YASA库可用，将使用YASA bandpower方法")
except ImportError:
    YASA_AVAILABLE = False
    logger.warning("YASA库不可用，将使用备用方法")


class YASABandpowerHelper:
    """
    YASA频段功率计算助手类
    
    专门用于计算频域比例，仅供前端显示使用
    """
    
    def __init__(self, sampling_rate: float = 512.0):
        """
        初始化YASA频段功率助手
        
        Args:
            sampling_rate: 采样率，默认512Hz
        """
        self.sampling_rate = sampling_rate
        self.yasa_available = YASA_AVAILABLE
        
        # 标准频段定义（与YASA默认定义一致）
        self.frequency_bands = {
            'delta': (0.5, 4.0),
            'theta': (4.0, 8.0), 
            'alpha': (8.0, 13.0),
            'beta': (13.0, 30.0),
            'gamma': (30.0, 40.0)
        }
        
        logger.info(f"YASABandpowerHelper初始化完成，采样率: {sampling_rate}Hz，YASA可用: {self.yasa_available}")
    
    def calculate_bandpower_ratios(self, 
                                   data: Union[np.ndarray, mne.io.Raw], 
                                   channel_name: str = 'EEG') -> Dict[str, float]:
        """
        计算频段功率比例
        
        Args:
            data: 输入数据，可以是numpy数组或MNE Raw对象
            channel_name: 通道名称，当输入为numpy数组时使用
            
        Returns:
            Dict[str, float]: 各频段的功率比例
        """
        try:
            if self.yasa_available:
                return self._yasa_bandpower_calculation(data, channel_name)
            else:
                return self._fallback_bandpower_calculation(data)
                
        except Exception as e:
            logger.error(f"频段功率计算失败: {str(e)}")
            return self._get_default_ratios()
    
    def _yasa_bandpower_calculation(self, 
                                    data: Union[np.ndarray, mne.io.Raw], 
                                    channel_name: str) -> Dict[str, float]:
        """
        使用YASA的bandpower方法计算频段功率
        
        Args:
            data: 输入数据
            channel_name: 通道名称
            
        Returns:
            Dict[str, float]: 各频段的功率比例
        """
        try:
            # 如果输入是numpy数组，转换为MNE Raw对象
            if isinstance(data, np.ndarray):
                raw = self._create_raw_from_array(data, channel_name)
            else:
                raw = data
            
            # 抑制警告
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                # 使用YASA的bandpower函数计算各频段功率
                band_powers = {}
                
                for band_name, (low_freq, high_freq) in self.frequency_bands.items():
                    # YASA bandpower函数计算指定频段的功率
                    power = yasa.bandpower(
                        raw, 
                        bands=[(low_freq, high_freq, band_name)],
                        ch_names=[channel_name] if channel_name in raw.ch_names else None,
                        relative=False  # 获取绝对功率，稍后手动计算比例
                    )
                    
                    # 提取功率值
                    if isinstance(power, pd.DataFrame):
                        band_powers[band_name] = float(power.iloc[0, 0]) if len(power) > 0 else 0.0
                    else:
                        band_powers[band_name] = float(power) if power is not None else 0.0
                
                # 计算总功率
                total_power = sum(band_powers.values())
                
                # 计算比例
                if total_power > 0:
                    power_ratios = {
                        band: round(power / total_power, 3) 
                        for band, power in band_powers.items()
                    }
                else:
                    power_ratios = self._get_default_ratios()
                
                logger.debug(f"YASA频段功率计算完成: {power_ratios}")
                return power_ratios
                
        except Exception as e:
            logger.error(f"YASA频段功率计算异常: {str(e)}")
            return self._fallback_bandpower_calculation(data)
    
    def _fallback_bandpower_calculation(self, data: Union[np.ndarray, mne.io.Raw]) -> Dict[str, float]:
        """
        备用频段功率计算方法（使用scipy）
        
        Args:
            data: 输入数据
            
        Returns:
            Dict[str, float]: 各频段的功率比例
        """
        try:
            from scipy import signal
            
            # 提取数据数组
            if isinstance(data, mne.io.Raw):
                data_array = data.get_data()[0]  # 取第一个通道
            else:
                data_array = data.flatten() if data.ndim > 1 else data
            
            # 计算功率谱密度
            freqs, psd = signal.welch(
                data_array, 
                fs=self.sampling_rate,
                nperseg=min(len(data_array), int(4 * self.sampling_rate))  # 4秒窗口
            )
            
            # 计算各频段功率
            band_powers = {}
            for band_name, (low_freq, high_freq) in self.frequency_bands.items():
                # 找到频段范围内的频率索引
                band_mask = (freqs >= low_freq) & (freqs <= high_freq)
                
                # 使用梯形积分计算频段功率
                if np.any(band_mask):
                    band_power = np.trapz(psd[band_mask], freqs[band_mask])
                else:
                    band_power = 0.0
                
                band_powers[band_name] = float(band_power)
            
            # 计算总功率和比例
            total_power = sum(band_powers.values())
            
            if total_power > 0:
                power_ratios = {
                    band: round(power / total_power, 3) 
                    for band, power in band_powers.items()
                }
            else:
                power_ratios = self._get_default_ratios()
            
            logger.debug(f"备用频段功率计算完成: {power_ratios}")
            return power_ratios
            
        except Exception as e:
            logger.error(f"备用频段功率计算异常: {str(e)}")
            return self._get_default_ratios()
    
    def _create_raw_from_array(self, data: np.ndarray, channel_name: str) -> mne.io.Raw:
        """
        从numpy数组创建MNE Raw对象
        
        Args:
            data: 数据数组
            channel_name: 通道名称
            
        Returns:
            mne.io.Raw: MNE Raw对象
        """
        # 确保数据是2D格式 (n_channels, n_samples)
        if data.ndim == 1:
            data = data.reshape(1, -1)
        
        # 创建info对象
        info = mne.create_info(
            ch_names=[channel_name],
            sfreq=self.sampling_rate,
            ch_types=['eeg']
        )
        
        # 创建Raw对象
        raw = mne.io.RawArray(data, info, verbose=False)
        return raw
    
    def _get_default_ratios(self) -> Dict[str, float]:
        """
        获取默认的频段功率比例
        
        Returns:
            Dict[str, float]: 默认比例（均匀分布）
        """
        num_bands = len(self.frequency_bands)
        default_ratio = round(1.0 / num_bands, 3)
        
        return {band: default_ratio for band in self.frequency_bands.keys()}
    
    def calculate_from_frequency_data(self, frequency_data) -> Dict[str, float]:
        """
        从FrequencyData对象计算频段比例（兼容现有数据结构）
        
        Args:
            frequency_data: FrequencyData对象
            
        Returns:
            Dict[str, float]: 频段功率比例
        """
        try:
            # 提取各频段功率
            delta = getattr(frequency_data, 'delta', 0.0)
            theta = getattr(frequency_data, 'theta', 0.0)
            low_alpha = getattr(frequency_data, 'low_alpha', 0.0)
            high_alpha = getattr(frequency_data, 'high_alpha', 0.0)
            low_beta = getattr(frequency_data, 'low_beta', 0.0)
            high_beta = getattr(frequency_data, 'high_beta', 0.0)
            low_gamma = getattr(frequency_data, 'low_gamma', 0.0)
            mid_gamma = getattr(frequency_data, 'mid_gamma', 0.0)
            
            # 合并频段
            alpha_total = low_alpha + high_alpha
            beta_total = low_beta + high_beta
            gamma_total = low_gamma + mid_gamma
            
            # 计算总功率
            total_power = delta + theta + alpha_total + beta_total + gamma_total
            
            if total_power == 0:
                return self._get_default_ratios()
            
            # 计算比例
            return {
                "delta": round(delta / total_power, 3),
                "theta": round(theta / total_power, 3),
                "alpha": round(alpha_total / total_power, 3),
                "beta": round(beta_total / total_power, 3),
                "gamma": round(gamma_total / total_power, 3)
            }
            
        except Exception as e:
            logger.error(f"从FrequencyData计算频段比例失败: {str(e)}")
            return self._get_default_ratios()
    
    def batch_calculate_ratios(self, data_list: List[Dict]) -> Dict[str, float]:
        """
        批量计算频段比例（用于多个数据点的平均）
        
        Args:
            data_list: 数据点列表
            
        Returns:
            Dict[str, float]: 平均频段功率比例
        """
        try:
            if not data_list:
                return self._get_default_ratios()
            
            # 累计各频段功率
            frequency_sums = {"delta": 0, "theta": 0, "alpha": 0, "beta": 0, "gamma": 0}
            valid_count = 0
            
            for data in data_list:
                frequency_data = data.get('frequency_data')
                if frequency_data:
                    ratios = self.calculate_from_frequency_data(frequency_data)
                    for band in frequency_sums:
                        frequency_sums[band] += ratios.get(band, 0)
                    valid_count += 1
            
            if valid_count == 0:
                return self._get_default_ratios()
            
            # 计算平均比例
            avg_ratios = {
                band: round(total / valid_count, 3) 
                for band, total in frequency_sums.items()
            }
            
            return avg_ratios
            
        except Exception as e:
            logger.error(f"批量计算频段比例失败: {str(e)}")
            return self._get_default_ratios()


# 全局实例，方便使用
yasa_bandpower_helper = YASABandpowerHelper()


def calculate_frequency_ratios(data: Union[np.ndarray, mne.io.Raw, object], 
                               channel_name: str = 'EEG') -> Dict[str, float]:
    """
    便捷函数：计算频段功率比例
    
    Args:
        data: 输入数据（numpy数组、MNE Raw对象或FrequencyData对象）
        channel_name: 通道名称
        
    Returns:
        Dict[str, float]: 频段功率比例
    """
    if hasattr(data, 'delta'):  # FrequencyData对象
        return yasa_bandpower_helper.calculate_from_frequency_data(data)
    else:  # numpy数组或MNE Raw对象
        return yasa_bandpower_helper.calculate_bandpower_ratios(data, channel_name)


def batch_calculate_frequency_ratios(data_list: List[Dict]) -> Dict[str, float]:
    """
    便捷函数：批量计算频段功率比例
    
    Args:
        data_list: 数据点列表
        
    Returns:
        Dict[str, float]: 平均频段功率比例
    """
    return yasa_bandpower_helper.batch_calculate_ratios(data_list)